#!/usr/bin/python2.5
# Copyright (C) 2007 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
  This package is an extension of transitfeed module.
  The goal of this package is to make a gtfs zip file editable.
  We will create temporary files for swapping and
  generate zip files as a result of editing.
"""

import hashlib

from transitfeed import *

# _TABLE_NAME is the table name of corresponding class in the database.
# The database stores data and intermediate changes which can be used to
# recover editing on the next run of the server and export changes to a
# new GTFS feed.
Fare._TABLE_NAME = 'fares'
FareRule._TABLE_NAME = 'fare_rules'
Transfer._TABLE_NAME = 'transfers'
Stop._TABLE_NAME = 'stops'
Shape._TABLE_NAME = 'shapes'

Agency._FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)',
                       'char(256)', 'char(256)', 'char(256)']

Route._FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)',
                      'integer', 'char(256)', 'char(256)',
                      'char(256)', 'char(256)', 'char(256)']
Trip._FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)',
                      'char(256)', 'char(256)', 'char(256)',
                      'char(256)']
Fare._FIELD_TYPES = ['char(256)', 'float', 'char(256)',
                      'integer', 'integer', 'integer']
FareRule._FIELD_TYPES = ['char(256)', 'char(256)', 'char(256)',
                      'char(256)', 'char(256)']
Stop._FIELD_TYPES = ['char(256)', 'char(256)', 'float',
                     'float', 'char(256)', 'char(256)',
                     'char(256)', 'char(256)', 'char(256)',
                     'char(256)']
Transfer._FIELD_TYPES = ['char(256)', 'char(256)', 'integer',
                     'integer']

Shape._FIELD_TYPES = ['char(256)', 'float', 'float', 'integer',
                     'float']

Agency._FEED_NAME = 'agency.txt'
Route._FEED_NAME = 'routes.txt'
Trip._FEED_NAME = 'trips.txt'
Fare._FEED_NAME = 'fare_attributes.txt'
FareRule._FEED_NAME = 'fare_rules.txt'
Stop._FEED_NAME = 'stops.txt'
Transfer._FEED_NAME = 'transfers.txt'
Shape._FEED_NAME = 'shapes.txt'

Agency._TYPE = 'agency'
Route._TYPE = 'route'
Trip._TYPE = 'trip'
Fare._TYPE = 'fare'
FareRule._TYPE = 'farerule'
Stop._TYPE = 'stop'
Transfer._TYPE = 'transfer'
Shape._TYPE = 'shape'

"""Utility functions"""
def TimeStringToSeconds(tm):
  if not tm:
    return None
  t = tm.strip().split(':')
  return int(t[0]) * 3600 + int(t[1]) * 60 + int(t[2])

def ReplaceBlankToNone(str):
  """Replace a string to None if it's a blank and return it."""
  if str == '':
    return None
  else:
    return str


class GTFSDatabase():
  """Initial version of GTFSDatabase, handling of database common
  functions
  TODO(wenxin): Add error handling for all database operation functions
  """

  def __init__(self, file='', memory_db=False):
    if memory_db:
      self._connection = sqlite.connect(':memory:')
    elif file:
      self._connection = sqlite.connect(file + '.db')
    else:
      try:
        self._temp_db_file = tempfile.NamedTemporaryFile()
        self._connection = sqlite.connect(self._temp_db_file.name)
      except sqlite.OperationalError:
        # Windows won't let a file be opened twice. mkstemp does not remove the
        # file when all handles to it are closed.
        self._temp_db_file = None
        (fd, self._temp_db_filename) = tempfile.mkstemp('.db')
        os.close(fd)
        self._connection = sqlite.connect(self._temp_db_filename)

    self._cursor = self._connection.cursor()
    self._file = file

  def Commit(self):
    """Shortcut to connection commit"""
    self._connection.commit()

  def DeleteStopTimeFromTrip(self, sequence, stopId, tripIds, deleteStop=False):
    """Given the stop_id, stop_sequence and trip_id, 
    delete stoptime from certain pattern."""
    for tripId in tripIds:
      sql = ('delete from `%s` where trip_id = %s and stop_id = "%s" '
             ' and stop_sequence = "%s";' % (
             EditableStopTime._TABLE_NAME,
             tripId,
             stopId,
             sequence)
             )
      self._cursor.execute(sql)
    if deleteStop:
      self.DeleteStop(stopId)

  def DeleteStop(self, stopId):
    """Delete stop from stop table by stopId"""
    if deleteStop:
      sql = ('delete from `%s` where stop_id = "%s";' %
             ( EditableStop._TABLE_NAME, stopId))
      self._cursor.execute(sql)

  def GetGTFSData(self, gtfsClass):
    """Get all gtfs data from table"""
    sql = """select %s from `%s`""" % (','.join(gtfsClass._FIELD_NAMES),
                                       gtfsClass._TABLE_NAME)
    self._cursor.execute(sql)
    return self._cursor.fetchall()

  def InsertTripStoptime(self, sequence, stopId, trips, arrivalTime='',
    departureTime='', isEnd=0):
    """Insert stoptime for trips in certain pattern"""
    for trip in trips:
      sql = ('update `%s` set stop_sequence = stop_sequence + 1 '
             'where stop_sequence >= "%s" and trip_id = "%s"' % (
              EditableStopTime._TABLE_NAME,
              sequence,
              trip.trip_id))
      self._cursor.execute(sql)

      # Arrival/departure time is required on generating svg graphics
      # and the starting point of the ending point must have this value
      # or the svg generator will fail to generate proper graph of
      # the pattern.
      # So as a temporary solution I copy the former points on the ends
      # as the new stop's default value.

      iTime = ''
      if arrivalTime or departureTime:
        if isEnd:
          iTime = trip.GetEndTimeString()
        else:
          iTime = trip.GetStartTimeString()

      sql = ('insert into `%s` (trip_id, arrival_time, departure_time,'
             'stop_id, stop_sequence) values("%s", "%s", "%s", "%s", %s);' % (
             EditableStopTime._TABLE_NAME,
             trip.trip_id,
             iTime,
             iTime,
             stopId,
             sequence)
             )
      self._cursor.execute(sql)


  def UpdateTripStoptime(self, sequence, stopId, tripIds):
    """Update stop id in stop time for certain trips"""
    sql = ('update `stop_times` set stop_id = "%s" where stop_sequence = "%s" '
           'and trip_id in (%s);' % (stopId, sequence, ','.join(tripIds)))
    self._cursor.execute(sql)

  def InsertStop(self, stop):
    """Insert a new Stop"""
    sql = ('insert into `%s` (`stop_id`, `stop_name`, `stop_lat`, '
           '`stop_lon`, `stop_desc`, `zone_id`, `stop_url`, `stop_code`, '
           '`location_type`, `parent_station`) values("%s", "%s", "%s", '
           '"%s","%s", "%s", "%s", "%s", "%s", "%s");' % (
              EditableStop._TABLE_NAME,
              stop.stop_id,
              stop.stop_name,
              stop.stop_lat,
              stop.stop_lon,
              stop.stop_desc,
              stop.zone_id,
              stop.stop_url,
              stop.stop_code,
              stop.location_type,
              stop.parent_station))
    self._cursor.execute(sql)

  def UpdateStopToStation(self, stop):
    """Update stop info"""
    sql = ('update %s set location_type = "1" where '
           'stop_id = "%s"' % (EditableStop._TABLE_NAME, stop.stop_id))
    self._cursor.execute(sql)

  def UpdateStop(self, stop):
    """Update stop info"""
    sql = ('update %s set stop_name = "%s", stop_lat = %s, stop_lon = %s'
           ' where stop_id = "%s"' % (EditableStop._TABLE_NAME,
                                      stop.stop_name,
                                      stop.stop_lat,
                                      stop.stop_lon,
                                      stop.stop_id))
    self._cursor.execute(sql)

  def UpdateShape(self, shape_id, points):
    """Update shape points with shape id"""
    idx = 1
    for point in points:
      row = [shape_id.encode('utf-8'), point[0], point[1], "%d" % idx, point[2]]
      self.UpdateTable(row, Shape._FIELD_NAMES, Shape._TABLE_NAME)
      idx += 1

  def DeleteShape(self, shape_id):
    """Delete Shape data by shape_id"""
    sql = ('delete from %s where shape_id = "%s"' %
           (Shape._TABLE_NAME, shape_id))
    self._cursor.execute(sql)
    return True

  def GetAll(self, table_name):
    """Get all from a table"""
    if self.IsTableExist(table_name):
      self._cursor.execute('select * from %s' % table_name)
      return self._cursor.fetchall()
    return []

  def GetStopTimesOfAll(self):
    """Get all stop times info of all trips from table stop_times"""
    sql = ("select trip_id, stop_id, stop_sequence, departure_time,"
           "arrival_time  from stop_times;")
    self._cursor.execute(sql)
    return self._cursor.fetchall()

  def GetStopsOfAll(self):
    """Get all stop info of all trips from table stops"""
    sql = 'select stop_id, stop_lat, stop_lon, stop_name from stops;'
    self._cursor.execute(sql)
    return self._cursor.fetchall()

  def GetTripRouteShapeRelation(self):
    """Get related info between trip, shape and route from table trips"""
    sql = 'select trip_id, shape_id, route_id from trips;'
    self._cursor.execute(sql)
    return self._cursor.fetchall()

  def GetStopTimesByTripId(self, header, trip_id):
    """Get all stoptimes"""
    self._cursor.execute('select ' + ','.join(header) + ' from stop_times '
                         'where trip_id=? order by stop_sequence',
                         (trip_id,))
    return self._cursor.fetchall()

  def GetStartOrEndTimeByTripId(self, trip_id, is_start=True):
    """Get start or end time fo a trip"""
    order = ''
    if not is_start:
      order = ' desc '
    self._cursor.execute(
        'select arrival_time, departure_time from stop_times where '
        'trip_id=? order by stop_sequence ' + order + ' limit 1',
        (trip_id,))
    return self._cursor.fetchone()

  def DuplicatedSequences(self, trip_id):
    """Get duplicated sequences as an array"""
    self._cursor.execute("select count(stop_sequence) as a from stop_times "
                         " where trip_id=? group by stop_sequence having "
                         "a > 1", (trip_id,))
    return self._cursor

  def IsGTFSInitialized(self):
    """Check if the database has been initialized correctly"""
    sql = ("select name from sqlite_master where type = 'table' and "
           "name = 'initialized'")
    h = self._cursor.execute(sql)
    ret = h.fetchall()
    return len(ret) == 1

  def SetGTFSAsInitialized(self):
    """Set the Db as correctly initialized by creating
    a table 'initialized' in the Db"""
    sql = 'create table `initialized` (tag text)'
    self._cursor.execute(sql)

  def IsTableExist(self, table_name):
    sql = ("select name from sqlite_master where type = 'table' and"
           " name = '%s'" % (table_name))
    h = self._cursor.execute(sql)
    ret = h.fetchall()
    return len(ret) == 1

  def FieldsExist(self, table_name, fields):
    sql = ("select sql from sqlite_master where type = "
           "'table' and tbl_name = '%s'" % (table_name))
    h = self._cursor.execute(sql)
    row = h.fetchone()
    old_fields = re.findall('\((.*)\)', row[0])[0].split(',')

    tmp_fields = []
    for row in old_fields:
      tmp_fields.append(row.strip().replace('`', '').split(' ')[0])
    for field in fields:
      if field not in tmp_fields:
        return False
    return True

  def FieldsTypesMatch(self, table_name, fields, types):
    sql = ("select sql from sqlite_master where type = "
           "'table' and tbl_name = '%s'" % (table_name))
    h = self._cursor.execute(sql)
    row = h.fetchone()
    old_fields = re.findall('\((.*)\)', row[0])[0].split(',')

    tmp_fields = {}
    for row in old_fields:
      tmp = row.strip().replace('`', '').split(' ')
      tmp_fields[tmp[0].upper()] = tmp[1].upper()
    i = 0
    for field in fields:
      if tmp_fields[field.upper()] != types[i].upper():
        return False
      i += 1
    return True

  def CreateTable(self, fields, types, name):
    """Create tables for gtfs data."""
    if self._cursor:
      str = 'drop table if exists %s;' % name
      self._cursor.execute(str)
      defs = []
      for i in range(0, len(fields)):
        defs.append(fields[i] + ' ' + types[i])
      str = 'create table %s (%s);' % (name, ','.join(defs))
      self._cursor.execute(str)
      self._connection.commit()

  def CreateIndex(self, fields, name):
    """Create an index for a table.
    
    Args:
      fields: list of field_name needs to create an index such as [field_name] 
              or [field_name1, field_name2]
      name: string of table name
    """
    if self._cursor:
      self._cursor.execute('CREATE INDEX %s_%s_index ON %s (%s);' % 
                           (name, ''.join(fields), name, ','.join(fields)))
      self.Commit()

  def UpdateTable(self, rows, header, name):
    """Update tables for gtfs data."""
    if self._cursor:
      tmp = []
      for row in rows:
        tmp.append(row.replace('"', '\''))
      str = ('insert into %s (%s) values("%s");' %
             (name, ', '.join(header), '", "'.join(tmp)))
      str = str.replace('"', '\"').replace('\\""', '\\"')
      self._cursor.execute(str)

  def InsertRows(self, rows, header, name):
    """Insert one or more rows from a list."""
    if self._cursor:
      str = ("""INSERT INTO %s (%s) VALUES (%s)""" %  
             (name, ','.join(header), ','.join(['?'] * len(header))))
      self._cursor.executemany(str, rows)

  def AddExtraColumns(self, fields, header, table):
    """Add extra fields not defined in _FIELD_NAMES."""
    for item in header:
      if item not in fields:
        str = 'alter table %s add `%s` char(256);' % (table, item)
        if self._cursor:
          self._cursor.execute(str)

  def UpdateShapeIdToTrip(self, schedule):
    """Update shape_id with trip.shape_id if shape_id is empty and 
    shapes has been generated before exporting feed."""
    if self._cursor:
      for trip in schedule.GetTripList():
        shape = schedule.GetShape(trip.shape_id)
        if trip.shape_id_hashed and shape and 'points' in dir(shape):
          self._cursor.execute('UPDATE trips set shape_id = "%s" WHERE '
                               'trip_id = "%s";' % 
                               (trip.shape_id, trip.trip_id))
      self.Commit()

class EditableGenericGTFSObject(GenericGTFSObject):
  """Base object of editable GTFS objects."""
  def __init__(self, field_dict=None):
    self._schedule = None
    if field_dict:
      for k, v in field_dict.iteritems():
        self.__dict__[k] = v

class Calendar(EditableGenericGTFSObject):
  """A class inherited from the EditableGenericGTFSObject class
  to enable editing on GTFS calendar.
  """
  _DAYS_OF_WEEK = ['monday', 'tuesday', 'wednesday', 'thursday',
                   'friday', 'saturday', 'sunday']
  _REQUIRED_FIELD_NAMES = ['service_id'] + _DAYS_OF_WEEK + [
                           'start_date', 'end_date']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES
  _TABLE_NAME = 'calendars'
  _FIELD_TYPES = ['char(256)', 'integer', 'integer',
                  'integer', 'integer', 'integer',
                  'integer', 'integer', 'char(256)',
                  'char(256)']
  _FEED_NAME = 'calendar.txt'
  _TYPE = 'calendar'

  def _IsValidDate(self, date):
    if re.match('^\d{8}$', date) == None:
      return False

    try:
      time.strptime(date, '%Y%m%d')
      return True
    except ValueError:
      return False

  def Validate(self, problems=default_problem_reporter):
    dates = {}
    for name in Calendar._REQUIRED_FIELD_NAMES:
      if IsEmpty(self.__dict__[name]):
        problems.MissingValue(name)
      elif name == 'start_date' or name == 'end_date':
        dates[name] = self.__dict__[name]
        if not self._IsValidDate(self.__dict__[name]):
          problems.InvalidValue(name, self.__dict__[name])
      elif name in Calendar._DAYS_OF_WEEK:
        value = self.__dict__[name]
        if (value != 0) and (value != 1):
          problems.InvalidValue(name, value)

    if ('start_date' in dir(dates) and
        'end_date' in dir(dates) and
        dates['end_date'] < dates['start_date']):
      problems.InvalidValue('end_date', dates['end_date'],
                            'end_date of %s is earlier than '
                            'start_date of "%s"' %
                            (dates['end_date'], dates['start_date']))

class CalendarDate(EditableGenericGTFSObject):
  """A class inherited from the EditableGenericGTFSObject class
  to enable editing on GTFS calendardate.
  """
  _REQUIRED_FIELD_NAMES = ['service_id', 'date', 'exception_type']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + []
  _TABLE_NAME = 'calendar_dates'
  _FIELD_TYPES = ['char(256)', 'char(256)', 'integer']
  _FEED_NAME = 'calendar_dates.txt'
  _TYPE = 'calendardate'

  def Validate(self, problems=default_problem_reporter):
    for name in CalendarDate._REQUIRED_FIELD_NAMES:
      if IsEmpty(self.__dict__[name]):
        problems.MissingValue(name)
      elif name == 'date':
        if not self._IsValidDate(self.__dict__[name]):
          problems.InvalidValue(name, self.__dict__[name])
      elif name == 'exception_type':
        value = self.__dict__[name]
        if (value != 1) and (value != 2):
          problems.InvalidValue(name, value)

class Frequency(EditableGenericGTFSObject):
  """A class inherited from the EditableGenericGTFSObject class
  to enable editing on GTFS frequency.
  """
  _REQUIRED_FIELD_NAMES = ['trip_id', 'start_time', 'end_time',
                           'headway_secs']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + []
  _TABLE_NAME = 'frequencies'
  _FIELD_TYPES = ['char(50)', 'char(20)', 'char(20)', 'integer']
  _FEED_NAME = 'frequencies.txt'
  _TYPE = 'frequency'

  def Validate(self, problems=default_problem_reporter):
    """place holder only"""
    pass

class EditableFare(Fare):
  def __init__(self, field_dict=None):
    self._schedule = None
    self.rules = []
    if field_dict:
      for k, v in field_dict.iteritems():
        self.__dict__[k] = v

class EditableFareRule(FareRule):
  def __init__(self, field_dict=None):
    self._schedule = None
    if field_dict:
      for k, v in field_dict.iteritems():
        self.__dict__[k] = v


class OriginalShape(Shape):
  """A class inherited from the Shape, to store original shape information
  before edited by the end user"""
  _TABLE_NAME = 'original_shapes'
  _SCHEDULE_REFERENCE = '_original_shapes'

class EditableShape(Shape):
  """A class inherited from class Shape to enable editing on GTFS Shape"""
  def Update(self, points, schedule=None):
    self.points = []
    for point in points:
      self.points.append(point.split(','))
    if schedule:
      db = schedule._db
      db.DeleteShape(self.shape_id)
      db.UpdateShape(self.shape_id, self.points)
      db.Commit()
    return True

Shape._SCHEDULE_REFERENCE = '_shapes'

class EditableStop(Stop):
  """Editable Stop inherited from Stop"""

  def __init__(self, field_dict=None):
    if not field_dict:
      field_dict = {'stop_id': '',
                    'stop_desc' : '',
                    'zone_id' : '',
                    'stop_url' : '',
                    'stop_code' : '',
                    'parent_station': '',
                    'stop_name': '',
                    'stop_lat': 0.0,
                    'stop_lon': 0.0,
                    'location_type': 0 }
    Stop.__init__(self, field_dict=field_dict)

  def Update(self, schedule):
    db = schedule._db
    db.UpdateStop(self)
    db.Commit()

  def ToStation(self, schedule):
    self.location_type = 1
    db = schedule._db
    db.UpdateStopToStation(self)
    db.Commit()

  def Insert(self, schedule):
    db = schedule._db
    db.InsertStop(self)
    db.Commit()

class EditableRoute(Route):
  def UpdatePattern(self, patternId, sequence, arrivalTime, departureTime,
    stopId, isEnd, schedule):
    """Update stoptime info of the specified pattern"""
    trips = []
    for trip in self._trips:
      if patternId == trip.pattern_id:
        trips.append(trip)
    db = schedule._db
    db.InsertTripStoptime(sequence, stopId, trips, arrivalTime, departureTime,
                          isEnd)
    db.Commit()

  def DeleteStopTimeFromPattern(self, patternId, sequence, stopId, schedule):
    """Delete stoptime info of the specified pattern"""
    tripIds = []
    for trip in self._trips:
      if patternId == trip.pattern_id:
        tripIds.append('"%s"' % trip.trip_id)
    db = schedule._db
    db.DeleteStopTimeFromTrip(sequence, stopId, tripIds)
    db.Commit()

class EditableStopTime(EditableGenericGTFSObject):
  """A class inherited from the EditableGenericGTFSObject class
  to enable editing on GTFS stoptime.
  """
  _TABLE_NAME = 'stop_times'
  _REQUIRED_FIELD_NAMES = ['trip_id', 'arrival_time',
                           'departure_time', 'stop_id',
                           'stop_sequence']
  _FIELD_NAMES = _REQUIRED_FIELD_NAMES + ['stop_headsign',
                  'pickup_type', 'drop_off_type',
                  'shape_dist_traveled']
  _FIELD_TYPES = ['char(50)', 'integer', 'integer',
                  'char(50)', 'integer', 'char(100)',
                  'integer', 'integer', 'float']
  _FEED_NAME = 'stop_times.txt'
  _TYPE = 'stoptime'

  def __getattr__(self, name):
    if name == 'stop_id':
      return self.stop.stop_id
    elif name == 'arrival_time':
      return (self.arrival_time != None and
          FormatSecondsSinceMidnight(self.arrival_time) or '')
    elif name == 'departure_time':
      return (self.departure_time != None and
          FormatSecondsSinceMidnight(self.departure_time) or '')
    elif name == 'shape_dist_traveled':
      return ''
    raise AttributeError(name)

  def GetTimeSecs(self):
    """Get Stop time info.

    Return the first of arrival_time and departure_time that is not None.
    If both are None return None.
    """
    if self.arrival_time != None:
      return self.arrival_time
    elif self.departure_time != None:
      return self.departure_time
    else:
      return None

class EditableTrip(Trip):
  """A class inherited from the Trip class to enable editing
  on GTFS feeds.
  """
  def Validate(self, problems, validate_children=True):
    """Validate attributes of this object.

    Check that this object has all required fields set to valid values without
    reference to the rest of the schedule. If the _schedule attribute is set
    then check that references such as route_id and service_id are correct.

    Args:
      problems: A ProblemReporter object
      validate_children: if True and the _schedule attribute is set then call
                         ValidateChildren
    """
    if IsEmpty(self.route_id):
      problems.MissingValue('route_id')
    if IsEmpty(self.service_id):
      problems.MissingValue('service_id')
    if IsEmpty(self.trip_id):
      problems.MissingValue('trip_id')
    if hasattr(self, 'direction_id') and (not IsEmpty(self.direction_id)) and \
        (self.direction_id != '0') and (self.direction_id != '1'):
      problems.InvalidValue('direction_id', self.direction_id,
                            'direction_id must be "0" or "1"')
    if self._schedule:
      if (self.shape_id and self.shape_id not in self._schedule._shapes and
          'shape_id_hashed' in dir(self) and not self.shape_id_hashed):
        problems.InvalidValue('shape_id', self.shape_id)
      if self.route_id and self.route_id not in self._schedule.routes:
        problems.InvalidValue('route_id', self.route_id)
      if (self.service_id and
          self.service_id not in self._schedule._calendars):
        problems.InvalidValue('service_id', self.service_id)

      if validate_children:
        self.ValidateChildren(problems)

  def GetStopTimes(self, problems=None):
    """Return a sorted list of StopTime objects for this trip.
    In theory problems=None should be safe because the data from
    database has been validated. See comment in _LoadStopTimes in
    module transitfeed for why this isn't always true."""
    stop_times = []
    db = self._schedule._db
    header = EditableStopTime._FIELD_NAMES
    rows = db.GetStopTimesByTripId(header, self.trip_id)
    for row in rows:
      stop = self._schedule.GetStop(row[3])
      stoptime = EditableStopTime(dict(zip(header, row)))
      if not stoptime.shape_dist_traveled:
        stoptime.shape_dist_traveled = '0.0'
      stoptime.shape_dist_traveled = float(
                                       stoptime.shape_dist_traveled)
      stoptime.__dict__['stop'] = stop
      stop_times.append(stoptime)
    return stop_times

  def GetTimeStops(self):
    """Return a list of (arrival_seconds, departure_seconds, stop)
    tuples.

    Caution: arrival_seconds and departure_seconds may be 0, a false
             value meaning a stop at midnight or None, it may also
             means the stop is untimed."""
    return [(TimeStringToSeconds(st.arrival_time),
             TimeStringToSeconds(st.departure_time),
             st.stop) for st in self.GetStopTimes()]

  def GetTimeStopsString(self):
    """Return a list of (arrival_time, departure_time, stop) tuples.

    Caution: arrival_time and departure_time may be 0, a false value
             meaning a stop at midnight or None, it may also means
             the stop is untimed."""
    return [(st.arrival_time, st.departure_time, st.stop) for st in
            self.GetStopTimes()]

  def GetStartTimeString(self, problems=default_problem_reporter):
    """Return the first time of the trip.
    TODO: For trips defined by frequency return the first time of
          the first trip."""
    db = self._schedule._db
    (arrival_time, departure_time) = db.GetStartOrEndTimeByTripId(self.trip_id)
    if arrival_time != None:
      return arrival_time
    elif departure_time != None:
      return departure_time
    else:
      problems.InvalidValue('departure_time', '',
                            'The first stop_time in trip %s is missing '
                            'times.' % self.trip_id)

  def GetStartTime(self, problems=default_problem_reporter):
    return TimeStringToSeconds(self.GetStartTimeString(problems))

  def GetEndTimeString(self, problems=default_problem_reporter):
    """Return the last time of the trip.
    TODO: For trips defined by frequency return the last time of the
          last trip."""
    db = self._schedule._db
    (arrival_time, departure_time) = db.GetStartOrEndTimeByTripId(self.trip_id,
                                                                  False)
    if departure_time != None:
      return departure_time
    elif arrival_time != None:
      return arrival_time
    else:
      problems.InvalidValue('arrival_time', '',
                            'The last stop_time in trip %s'
                            'is missing times.' % self.trip_id)

  def GetEndTime(self, problems=default_problem_reporter):
    return TimeStringToSeconds(self.GetEndTimeString(problems))

  def ValidateChildren(self, problems=None):
    """Validate StopTimes and headways of this trip."""
    assert self._schedule, """Trip must be in a schedule to
            ValidateChildren"""
    # TODO: validate distance values in stop times (if applicable)
    db = self._schedule._db
    for row in db.DuplicatedSequences(self.trip_id):
      problems.InvalidValue('stop_sequence', row[0],
                            'Duplicate stop_sequence in trip_id %s' %
                            self.trip_id)

    stoptimes = self.GetStopTimes(problems)
    if stoptimes:
      if (stoptimes[0].arrival_time is None and
          stoptimes[0].departure_time is None):
        problems.OtherProblem(
          'No time for start of trip_id "%s""' % (self.trip_id))
      if (stoptimes[-1].arrival_time is None and
          stoptimes[-1].departure_time is None):
        problems.OtherProblem(
          'No time for end of trip_id "%s""' % (self.trip_id))

      # Sorts the stoptimes by sequence and then checks that the arrival time
      # for each time point is after the departure time of the previous.
      stoptimes.sort(key=lambda x: x.stop_sequence)
      prev_departure = 0
      prev_stop = None
      prev_distance = None
      try:
        route_type = self._schedule.GetRoute(
                                      self.route_id).route_type
        max_speed = Route._ROUTE_TYPES[route_type]['max_speed']
      except KeyError, e:
        # If route_type cannot be found, assume it is 0 (Tram) for checking
        # speeds between stops.
        max_speed = Route._ROUTE_TYPES[0]['max_speed']
      for timepoint in stoptimes:
        # Distance should be a nonnegative float number, so it should be
        # always larger than None.

        distance = timepoint.shape_dist_traveled
        if not distance == 0 and distance is not None:
          if distance > prev_distance and distance >= 0:
            prev_distance = distance
          else:
            problems.InvalidValue('stoptimes.shape_dist_traveled',
                                  distance,
              'For the trip %s the stop %s has shape_dist_travled=%f, '
              'which should be larger than the previous ones. In this '
              'case, the previous distance was %s.' %
              (self.trip_id,
               timepoint.stop_id,
               distance,
               prev_distance))

        if timepoint.arrival_time is not None:
          self._CheckSpeed(prev_stop,
                           timepoint.stop,
                           prev_departure,
                           timepoint.arrival_time,
                           max_speed,
                           problems)

          if timepoint.arrival_time >= prev_departure:
            prev_departure = timepoint.departure_time
            prev_stop = timepoint.stop
          else:
            problems.OtherProblem('Timetravel detected! Arrival time'
                                  ' is before previous departure '
                                  'at sequence number %s in trip %s'
                                  % (timepoint.stop_sequence,
                                     self.trip_id))

      if self.shape_id and self.shape_id in self._schedule._shapes:
        shape = self._schedule.GetShape(self.shape_id)
        max_shape_dist = shape.max_distance
        st = stoptimes[-1]
        if (st.shape_dist_traveled and
            st.shape_dist_traveled > max_shape_dist):
          problems.OtherProblem(
              'In stop_times.txt, the stop with trip_id=%s and '
              'stop_sequence=%d has shape_dist_traveled=%f, '
              'which is larger than the max shape_dist_traveled=%f'
              ' of the corresponding shape (shape_id=%s)' %
              (self.trip_id, st.stop_sequence,
               st.shape_dist_traveled, max_shape_dist,
               self.shape_id), type=TYPE_WARNING)

        # shape_dist_traveled is valid in shape if max_shape_dist larger than
        # 0.
        if max_shape_dist > 0:
          for st in stoptimes:
            if st.shape_dist_traveled is None:
              continue
            pt = shape.GetPointWithDistanceTraveled(
                        st.shape_dist_traveled)
            if pt:
              stop = self._schedule.GetStop(st.stop_id)
              distance = ApproximateDistance(stop.stop_lat,
                                             stop.stop_lon,
                                             pt[0], pt[1])
              if distance > MAX_DISTANCE_FROM_STOP_TO_SHAPE:
                problems.StopTooFarFromShapeWithDistTraveled(
                    self.trip_id, stop.stop_name, stop.stop_id,
                    pt[2], self.shape_id, distance,
                    MAX_DISTANCE_FROM_STOP_TO_SHAPE)

    # O(n^2), but we don't anticipate many headway periods per trip
    for headway_index, headway in enumerate(self._headways[0:-1]):
      for other in self._headways[headway_index + 1:]:
        if (other[0] < headway[1]) and (other[1] > headway[0]):
          problems.OtherProblem('Trip contains overlapping headway '
                                'periods %s and %s' %
                                (self._HeadwayOutputTuple(headway),
                                 self._HeadwayOutputTuple(other)))

class EditableSchedule(Schedule):
  """A class inherited from the Schedule class to enable editing
  on GTFS feeds.
  """

  def __init__(self, feed_path=None,
               problem_reporter=default_problem_reporter,
               memory_db=False, check_duplicate_trips=False):
    """an overriding on the parent __init__ function.
    add _original_shapes"""
    self._feed_path = feed_path
    Schedule.__init__(self, problem_reporter, memory_db,
                      check_duplicate_trips)
    self._original_shapes = {}  # Original shape_id to Shape

  def GetShape(self, shape_id):
    if shape_id not in self._shapes:
      return None
    return self._shapes[shape_id]

  def SaveToFeedFile(self):
    """Save database info into a zip file"""
    basename, extension = os.path.splitext(self._feed_path)
    self.WriteGoogleTransitFeed(basename + "-new" + extension)    

  def ExportGTFSDataToTextFile(self, gtfsClass, archive):
    """Export a gtfsClass data into an archive"""
    str = StringIO.StringIO()
    writer = CsvUnicodeWriter(str)
    columns = gtfsClass._FIELD_NAMES
    writer.writerow(columns)
    writer.writerows(self._db.GetGTFSData(gtfsClass))
    self._WriteArchiveString(archive, gtfsClass._FEED_NAME, str)


  def WriteGoogleTransitFeed(self, file):
    """Output this schedule as a Google Transit Feed Data in file """
    archive = zipfile.ZipFile(file, 'w')
    classes = [
      Agency, EditableStop, EditableRoute,
      Calendar, CalendarDate, EditableStopTime,
      EditableTrip, EditableShape, Frequency,
      EditableFare, EditableFareRule, Transfer]
    for gtfsClass in classes:
      if gtfsClass == EditableTrip:
        self._db.UpdateShapeIdToTrip(self)
      self.ExportGTFSDataToTextFile(gtfsClass, archive)
    archive.close()

  def ConnectDb(self, memory_db):
    """Overridden to enable path specified db"""
    self._db = GTFSDatabase(self._feed_path, memory_db=memory_db)
    self._connection = self._db._connection
    self._cursor = self._db._connection.cursor()

  def Load(self, feed_path, extra_validation=False):
    """An overriding on the parent Load. Check the import
    and add flag for it."""
    self._imported = self._db.IsGTFSInitialized()

    loader = EditableLoader(feed_path,
                            self,
                            problems=self.problem_reporter,
                            extra_validation=extra_validation,
                            imported=self._imported,
                            connection=self._connection,
                            db=self._db)
    loader.Load()

    if not self._imported:
      self._db.SetGTFSAsInitialized()

  def AddGTFSObject(self, obj, obj_id, schedule_obj,
                    problem_reporter=None):
    if not problem_reporter:
      problem_reporter = self.problem_reporter

    obj.Validate(problem_reporter)

    if schedule_obj not in dir(self):
      self.__dict__[schedule_obj] = {}

    if obj.__dict__[obj_id] in self.__dict__[schedule_obj]:
      problem_reporter.DuplicateID(obj_id, obj.__dict__[obj_id])
      return
    self.__dict__[schedule_obj][obj.__dict__[obj_id]] = obj

  def AddGTFSNoneUniqueObject(self, obj, obj_id, schedule_obj,
                              problem_reporter=None):
    if not problem_reporter:
      problem_reporter = self.problem_reporter

    obj.Validate(problem_reporter)

    if schedule_obj not in dir(self):
      self.__dict__[schedule_obj] = {}
    self.__dict__[schedule_obj].setdefault(obj.__dict__[obj_id],
                                           []).append(obj)

  def AddObject(self, obj, problem_reporter=None):
    """add to Schedule by obj's _TYPE"""
    if obj._TYPE == 'agency':
      self.AddAgencyObject(obj, problem_reporter)
    if obj._TYPE == 'route':
      self.AddRouteObject(obj, problem_reporter)
    if obj._TYPE == 'trip':
      self.AddTripObject(obj, problem_reporter)
    if obj._TYPE == 'stop':
      self.AddStopObject(obj, problem_reporter)
    if obj._TYPE == 'transfer':
      self.AddTransferObject(obj, problem_reporter)
    if obj._TYPE == 'fare':
      self.AddFareObject(obj, problem_reporter)
    if obj._TYPE == 'farerule':
      self.AddFareRuleObject(obj, problem_reporter)
    if obj._TYPE == 'calendar':
      self.AddGTFSObject(obj, 'service_id', '_calendars', problem_reporter)
    if obj._TYPE == 'frequency':
      self.AddGTFSNoneUniqueObject(obj, 'trip_id', '_frequencies',
                                   problem_reporter)


class EditableLoader(Loader):
  """Inherited from the Loader to enable editing on GTFS feeds."""
  def __init__(self,
               feed_path=None,
               schedule=None,
               problems=default_problem_reporter,
               extra_validation=False,
               load_stop_times=True,
               memory_db=True,
               zip=None,
               check_duplicate_trips=False,
               imported=False,
               connection=None,
               db=None):
    self._imported = imported
    if connection:
      self._connection = connection
      self._db = db
      self._cursor = self._db._cursor

    Loader.__init__(self, feed_path, schedule, problems, extra_validation, 
                    load_stop_times, memory_db, zip, check_duplicate_trips)

  def _ImportDataFromCsvToDb(self, gtfs_obj):
    """Import data from a file in gtfs feed into a table in database
    Args:
      gtfs_obj: a class drived from GenericGTFSObject
    """
    extra_columns_added = False
    self._db.CreateTable(gtfs_obj._FIELD_NAMES,
                         gtfs_obj._FIELD_TYPES,
                         gtfs_obj._TABLE_NAME)

    if not self._HasFile(gtfs_obj._FEED_NAME):
      return

    # Create index on table stop_times and trips which can improve the query
    # speed.
    if gtfs_obj._TABLE_NAME in ('stop_times', 'trips'):
      self._db.CreateIndex(['trip_id'], gtfs_obj._TABLE_NAME)

    rows = []
    header = []
    for (d, row_num, header, row) in \
     self._ReadCsvDict(gtfs_obj._FEED_NAME,
                       gtfs_obj._FIELD_NAMES,
                       gtfs_obj._REQUIRED_FIELD_NAMES):
      if not extra_columns_added:
        columns = self._db.AddExtraColumns(gtfs_obj._FIELD_NAMES,
                                           header,
                                           gtfs_obj._TABLE_NAME)
        extra_columns_added = True
      rows.append(row)

    # Store all the listings in rows and insert into DB with executemany(), 
    # this can reduce 1/3 time.
    if rows:
      self._db.InsertRows(rows, header, gtfs_obj._TABLE_NAME)
    del rows

  def _LoadObjectsFromDb(self, gtfs_obj):
    """Load gtfs objects from database into schedule"""
    header = gtfs_obj._FIELD_NAMES
    rows = self._db.GetAll(gtfs_obj._TABLE_NAME)
    row_num = 1
    for row in rows:
      # Convert '' to None because problems will throw errors if field is not 
      # None but '' in the objects.   
      formatted_row = tuple(ReplaceBlankToNone(r) for r in row)
      d = dict(zip(header, formatted_row))
      self._problems.SetFileContext(gtfs_obj._FEED_NAME,
                                    row_num,
                                    formatted_row,
                                    header)
      gtfs = gtfs_obj(field_dict=d)
      self._schedule.AddObject(gtfs, self._problems)
      self._problems.ClearContext()
      row_num += 1

  def _LoadObject(self, gtfs_obj, needLoad=True):
    """Always initialize from database, and import data into
    database if it is not imported"""
    if '_imported' not in dir(self) or not self._imported:
      self._ImportDataFromCsvToDb(gtfs_obj)
    if needLoad:
      self._LoadObjectsFromDb(gtfs_obj)

  def _BindTripShapeStopTimesInfo(self):
    """initialize relationship between entities"""

    # Self.GetStopTimesOfAll() returns an array in format of
    # (trip_id, stop_id, stop_sequence, departure_time, arrival_time)
    stoptimes_of_all = self._db.GetStopTimesOfAll()
    stoptimes_in_trip = {}
    for stoptime in stoptimes_of_all:
      stoptimes_in_trip.setdefault(stoptime[0], []).append([stoptime[1],
                                                            stoptime[2],
                                                            stoptime[3],
                                                            stoptime[4]])
    for trip_id in stoptimes_in_trip:
      def sort_stoptimes(t1, t2):
        return cmp(t1[2], t2[2])
      stoptimes_in_trip[trip_id].sort(sort_stoptimes)

    # Self._db.GetStopsOfAll() returns an array in format of
    # (stop_id, stop_lat, stop_lon, stop_name)
    stops_of_all = self._db.GetStopsOfAll()
    stop_dict = {}
    for stop in stops_of_all:
      stop_dict.setdefault(stop[0], []).append(stop)

    hashed_shapes = {}
    full_info_stop_lists = {}

    # Self._db.GetTripRouteShapeRelation() returns an array in format of
    # (trip_id, shape_id, route_id)
    trip_shape_route_relation = self._db.GetTripRouteShapeRelation()
    for relation in trip_shape_route_relation:
      trip = self._schedule.GetTrip(relation[0])
      
      # Generate trip.pattern_id with the hash value of a tuple of trip's stops.
      # It's the same with Trip.__getattr__('pattern_id') in _transitfeed.py 
      # and it will cost some time. The trip.pattern_id can be generated 
      # before starting the web server so that pattern_id needs not be generated
      # when a route is selected and patterns of the route are shown.
      trip.pattern_id = hash(tuple(self._schedule.GetStop(stop_time[0]) \
                             for stop_time in stoptimes_in_trip[relation[0]]))

      str = ''.join(stoptime[0] for stoptime in stoptimes_in_trip[relation[0]])
      m = hashlib.md5()
      m.update(str + relation[2])
      shape_hash = m.hexdigest()
      if not shape_hash in hashed_shapes:
        hashed_shapes[shape_hash] = False

      if not hashed_shapes[shape_hash]:
        full_stop_info_list = []
        for stoptime in stoptimes_in_trip[relation[0]]:
          stop = stop_dict[stoptime[0]][0] # Get first stop
          full_stop_info = [stoptime[1],
                            stop[1],
                            stop[2],
                            stop[3],
                            stoptime[3]]
          full_stop_info_list.append(full_stop_info)
        full_info_stop_lists[shape_hash] = full_stop_info_list
        hashed_shapes[shape_hash] = True
      if not relation[1]:
        trip.shape_id = shape_hash
        trip.shape_id_hashed = True
      else:
        trip.shape_id = relation[1]
        trip.shape_id_hashed = False
      full_stop_info_list = full_info_stop_lists[shape_hash]
      trip.shape_desc = '%s to %s, %d stops' % (full_stop_info_list[0][3],
                                                full_stop_info_list[-1][3],
                                                len(full_stop_info_list))
      trip.start_time = stoptimes_in_trip[relation[0]][0][3]

  def _LoadShapesIntoObjects(self, shape_class):
    header = shape_class._FIELD_NAMES
    rows = self._db.GetAll(shape_class._TABLE_NAME)
    row_num = 1
    shapes = {}
    for row in rows:
      d = dict(zip(header, row))
      file_context = (shape_class._FEED_NAME, row_num, row, header)
      self._problems.SetFileContext(*file_context)
      (shape_id, lat, lon, seq, dist) = row
      shapes.setdefault(shape_id, []).append((seq, lat, lon, dist, file_context))
      self._problems.ClearContext()
      row_num += 1

    for shape_id, points in shapes.items():
      shape = shape_class(shape_id)
      points.sort()
      last_seq = 1
      for (seq, lat, lon, dist, file_context) in points:
        if (seq != last_seq):
          self._problems.SetFileContext(*file_context)
          self._problems.InvalidValue('shape_pt_sequence', seq,
                                      'The sequence value is not equal to the'
                                      ' anticipated value %d of shape %s' %
                                      (last_seq, shape_id))
        last_seq = seq + 1
        shape.AddPoint(lat, lon, dist, self._problems)
        self._problems.ClearContext()
      self._schedule.AddGTFSObject(shape, 'shape_id',
                                   shape_class._SCHEDULE_REFERENCE,
                                   self._problems)

  def Load(self):
    self._problems.ClearContext()
    if not self._DetermineFormat():
      return self._schedule
    
    self._CheckFileNames()

    self._LoadObject(Agency)
    self._LoadObject(EditableStop)
    self._LoadObject(EditableRoute)
    self._LoadObject(Calendar)
    self._LoadObject(CalendarDate)

    self._LoadObject(EditableStopTime, False)
    self._LoadObject(EditableTrip)

    self._LoadObject(EditableShape, False)
    self._LoadShapesIntoObjects(EditableShape)

    self._LoadObject(OriginalShape, False)
    self._LoadShapesIntoObjects(OriginalShape)

    self._BindTripShapeStopTimesInfo()

    self._LoadObject(Frequency)

    self._LoadObject(EditableFare)
    self._LoadObject(EditableFareRule)
    self._LoadObject(Transfer)

    if self._zip:
      self._zip.close()
      self._zip = None

    if self._extra_validation:
      self._schedule.Validate(self._problems, validate_children=False)

    return self._schedule

